import com.aweme.harmony.mindspore.llm.jni.MSLiteLLMInferNativeKMP
import kotlinx.cinterop.CPointer

import com.ugc.harmony.clientai.lang.model.engine.LangEngineCallback
import com.ugc.harmony.clientai.lang.model.engine.LangEnginePerfCallback
import kotlinx.cinterop.ByteVar
import kotlinx.cinterop.COpaquePointer
import kotlinx.cinterop.StableRef
import kotlinx.cinterop.allocPointerTo
import kotlinx.cinterop.asStableRef
import kotlinx.cinterop.memScoped
import kotlinx.cinterop.ptr
import kotlinx.cinterop.staticCFunction
import kotlinx.cinterop.toKString
import kotlinx.cinterop.value
import platform.posix.usleep
import kotlin.native.concurrent.TransferMode
import kotlin.native.concurrent.Worker


/**
 * LLM推理器高级API封装
 * 负责模型生命周期管理和文本生成任务，对Kotlin调用者暴露友好的Map接口
 */
class MSLiteLLMInfer(modelPath: String, config: Map<String, String> = emptyMap(), isRecordPerfData: Boolean = false) {

    private var isInitialized = false
    private var handle: CPointer<*>? = null
    private var isGenerating = false
    private var currentCallbackStableRef: StableRef<LangEngineCallback>? = null

    init {
        handle = MSLiteLLMInferNativeKMP.LLMInfer_Native_Create_Kmp()
        initialize(modelPath, config, isRecordPerfData);
    }

    /**
     * 初始化模型
     * @param modelPath 模型文件路径
     * @param config 配置参数Map
     */
    private fun initialize(modelPath: String, config: Map<String, String>, isRecordPerfData: Boolean = false) {
        val success = MSLiteLLMInferNativeKMP.LLMInfer_Native_Build_Kmp(
            handle,
            modelPath,
            convertMapToNativeString(config),
            isRecordPerfData
        )
        isInitialized = success == 0
        if (success != 0) {
            dispose();
            throw IllegalStateException("模型初始化失败: $modelPath")
        }
    }

    /**
     * 生成文本 - 主要接口
     * @param input 输入文本
     * @return 生成的文本结果
     */
    fun generate(input: String): String {
        require(isInitialized) { "模型未初始化，请先调用initialize()方法" }
        return memScoped {
            val outputTextPtr = allocPointerTo<ByteVar>()
            val result = MSLiteLLMInferNativeKMP.LLMInfer_Native_Generate_Kmp(handle, input, outputTextPtr.ptr)

            if (result != 0) {
                // 可以根据不同的result返回不同的错误提示
                return@memScoped "(生成失败，错误码: $result)"
            }

            val outputCString: CPointer<ByteVar>? = outputTextPtr.value
            return@memScoped outputCString?.toKString() ?: "(生成结果为空)"
        }
    }

    /**
     * 生成文本 - 主要接口
     * @param input 输入文本
     * @param callback callback fun
     * @return 生成的文本结果
     */
    fun generate(input: String, callback: LangEnginePerfCallback): String {
        require(isInitialized) { "模型未初始化，请先调用initialize()方法" }
        return memScoped { // memScoped 块用于自动管理分配的内存在块结束时释放[1](@ref)
            // 1. 为输出参数 char** output_text 分配指针
            val outputTextPtr = allocPointerTo<ByteVar>()
            // 1. 创建 StableRef 来保护 callback 对象，防止被 GC
            val stableRef: StableRef<LangEnginePerfCallback>? = callback?.let {
                StableRef.create(it)
            }
            // 将 StableRef 作为 COpaquePointer 传递
            val callbackKotlinPtr: COpaquePointer? = stableRef?.asCPointer()

            // 2. 创建 C 兼容的回调函数指针
            val cCallback = staticCFunction { result: CPointer<ByteVar>?, data: COpaquePointer? ->
                // 2.1 将传入的 data 指针（即 StableRef 的指针）转换回 StableRef
                val ref: StableRef<LangEnginePerfCallback>? = data?.asStableRef()
                // 2.2 从 StableRef 中获取原始的 Kotlin callback 对象
                val targetCallback: LangEnginePerfCallback? = ref?.get()

                // 2.3 执行回调逻辑
                val kotlinString = result?.toKString() ?: ""
                println("Received from C: $kotlinString")
                targetCallback?.onPerformance(kotlinString) // 现在可以安全调用

                // 2.4 (可选) 如果回调是单次的，在此处释放 StableRef
                // ref?.dispose()
                Unit
            }

            // 3. 调用 C 函数
            val resultCode = MSLiteLLMInferNativeKMP.LLMInfer_Native_Generate_With_CallBack_Kmp(
                handle,
                input,
                outputTextPtr.ptr,
                cCallback,
                callbackKotlinPtr
            )

            if (resultCode != 0) {
                // 可以根据不同的result返回不同的错误提示
                return@memScoped "(生成失败，错误码: $resultCode)"
            }

            val outputCString: CPointer<ByteVar>? = outputTextPtr.value
            return@memScoped outputCString?.toKString() ?: "(生成结果为空)"
        }
    }


    fun generateAsync(input: String, callback: LangEngineCallback) {
        require(!isGenerating) { "已有异步操作正在进行" }

        isGenerating = true
        // 创建 StableRef 防止回调对象被 GC
        currentCallbackStableRef = StableRef.create(callback)

        memScoped {
            val result = MSLiteLLMInferNativeKMP.LLMInfer_Native_Generate_Async_Kmp(handle, input)
            if (result != 0) {
                cleanup()
                callback.onError("启动异步生成失败，错误码: $result")
                return
            }

            // 启动后台线程轮询状态
            startPollingThread()
        }
    }

    private fun startPollingThread() {
        val job = Worker.start()
        println("worker start\n")
        // 提取需要传递的状态
        val localHandle = this.handle
        val localIsGenerating = this::isGenerating
        val localCurrentCallbackStableRef = this::currentCallbackStableRef

        job.execute(TransferMode.SAFE, {
            Triple(localHandle, localIsGenerating, localCurrentCallbackStableRef)
        }) { triple ->
            val (handle, isGeneratingRef, currentCallbackStableRefRef) = triple
            var loopCnt = 0
            while (isGeneratingRef.get()) {
                memScoped {
                    val isEnd = MSLiteLLMInferNativeKMP.LLMInfer_Native_Generate_Is_End_Kmp(handle)
                    println("loop $loopCnt \n")
                    loopCnt++
                    if (isEnd) {
                        val outputTextPtr = allocPointerTo<ByteVar>()
                        val resultCode = MSLiteLLMInferNativeKMP.LLMInfer_Native_Get_Generate_Async_Result_Kmp(
                            handle,
                            outputTextPtr.ptr
                        )

                        if (resultCode == 0) {
                            val output = outputTextPtr.value?.toKString() ?: ""
                            currentCallbackStableRefRef.get()?.let { stableRef ->
                                try {
                                    stableRef.get().onReceive(output, true)
                                } catch (e: Exception) {
                                    println("回调执行异常: ${e.message}")
                                }
                            }
                        } else {
                            currentCallbackStableRefRef.get()?.let { stableRef ->
                                try {
                                    stableRef.get().onError("异步生成失败，错误码: $resultCode")
                                } catch (e: Exception) {
                                    println("回调执行异常: ${e.message}")
                                }
                            }
                        }

                        // 清理状态
                        isGeneratingRef.set(false)
                        currentCallbackStableRefRef.get()?.dispose()
                        currentCallbackStableRefRef.set(null)
                        break
                    }
                }

                usleep(100000u) // 休眠100ms避免过度轮询
            }
        }
    }

    fun cleanup() {
        isGenerating = false
        currentCallbackStableRef?.dispose()
        currentCallbackStableRef = null
    }

    fun isBusy(): Boolean = isGenerating

    /**
     * 释放模型资源
     */
    fun dispose() {
        if (isInitialized) {
            MSLiteLLMInferNativeKMP.LLMInfer_Native_Destroy_Kmp(handle)
            isInitialized = false
        }
    }

    /**
     * 释放模型资源
     */
    fun stopGenerating() {
        if (isInitialized) {
            MSLiteLLMInferNativeKMP.LLMInfer_Native_Stop_Generate_Kmp(handle)
        }
    }

    /**
     * 检查模型状态
     */
    fun isInitialized(): Boolean = isInitialized

    protected fun finalize() {
        dispose()
    }

    // region 工具方法

    /**
     * 将Map转换为"key1:value1;key2:value2;..."格式的字符串
     */
    private fun convertMapToNativeString(map: Map<String, String>): String {
        return map.entries.joinToString(";") { (key, value) ->
            "$key:$value"
        }
    }
}